import torch
import torchvision.models as models
from torchvision.transforms import Compose, ToTensor, Normalize
import pandas as pd


# --------------------- 模型加载 ---------------------
def load_model(model_name, model_path):
    """
    加载训练好的模型（LeNet/AlexNet/VGGNet）
    :param model_name: 模型名称，可选值: lenet, alexnet, vggnet
    :param model_path: 模型文件路径（如 models/lenet.pth）
    :return: 加载好的 PyTorch 模型（eval 模式）
    """
    # 1. 初始化模型（根据实际训练代码调整模型定义！！！）
    if model_name == "LeNet_mammals_best_model":
        # 需与训练时的 LeNet 结构一致，这里示例最简版
        model = models.LeNet(weights=None)  # weights=None 避免自动下载官方权重
    elif model_name == "alex_mammals_best_model":
        model = models.alexnet(weights=None)
    elif model_name == "vgg16_mammals_best_model":
        model = models.vgg16(weights=None)  # 以 VGG-16 为例，需与训练一致
    else:
        raise ValueError("不支持的模型: {}".format(model_name))

    # 2. 加载训练好的权重
    model.load_state_dict(torch.load(model_path, map_location="cuda"))
    model.eval()  # 推理模式
    return model

# --------------------- 图像预处理 ---------------------
def preprocess_image(image):
    """
    图像预处理：转为 Tensor + 归一化
    :param image: PIL Image 对象
    :return: 预处理后的 Tensor (1, 3, H, W)
    """
    transform = Compose([
        ToTensor(),  # 转为 [0,1] Tensor
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # ImageNet 归一化
    ])
    return transform(image).unsqueeze(0)  # 增加 batch 维度

# --------------------- 模型推理 ---------------------
def predict(model, image_tensor):
    """
    模型推理：输出 top-5 预测结果
    :param model: 加载好的 PyTorch 模型
    :param image_tensor: 预处理后的图像 Tensor
    :return: top-5 类别索引、top-5 概率
    """
    with torch.no_grad():  # 关闭梯度计算
        outputs = model(image_tensor)
        probabilities = torch.softmax(outputs, dim=1).squeeze()  # 计算概率
        top5_probs, top5_indices = torch.topk(probabilities, 5)  # 取 top-5
    return top5_indices.numpy(), top5_probs.numpy()

# --------------------- 模型参数计算 ---------------------
def count_parameters(model):
    """
    计算模型总参数量
    :param model: PyTorch 模型
    :return: 总参数量（单位：百万）
    """
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return total_params / 1e6  # 转为百万级，方便展示

# --------------------- 模型结构对比表 ---------------------
def get_model_info_table():
    """
    返回 LeNet/AlexNet/VGGNet 的结构对比表格（DataFrame）
    """
    data = {
        "模型": ["LeNet", "AlexNet", "VGGNet"],
        "卷积层特点": ["小卷积核 (5x5)", "大尺寸输入 (227x227)", "堆叠小卷积核 (3x3)"],
        "池化层": ["平均池化", "最大池化", "多层最大池化"],
        "经典应用": ["手写数字识别", "ImageNet 分类", "图像特征提取"]
    }
    return pd.DataFrame(data)